
# UNet for Left Atrium Segmentation

## Directory Structure
```
UNet_project/
├── model.py             # U-Net model architecture
├── data.py              # Data loading and preprocessing
├── train.py             # Training loop with weighted focal loss
├── test.py              # Evaluation script with Dice and HD95
├── loss_metrics.py      # Metrics computation and visualization
├── training_config.py   # Augmentation and callbacks configuration
├── train.txt            # Training set list
├── test.txt             # Test set list
```

## How to Run

### Training
```bash
python train.py
```

### Evaluation
```bash
python test.py
```

## Datasets

This study uses the open-source **ImageCHD** dataset to evaluate the performance of the U-Net architecture for **Left Atrium (LA)** segmentation.  
🔗 [Kaggle Dataset](https://www.kaggle.com/datasets/xiaoweixumedicalai/imagechd)

- 2D axial slices are extracted from 3D `.nii.gz` volumes.
- Class label 3 is used for segmenting the **Left Atrium (LA)**.
- All slices are resized to **512 × 512 pixels**.
- The in-plane resolution is **0.25 × 0.25 mm**.
- The **95th percentile Hausdorff Distance (HD95)** is calculated in pixels and converted to millimeters by multiplying by `0.25`.

## Model Details

- Based on the original **U-Net** architecture with 5 encoder-decoder stages.
- Includes skip connections for spatial detail preservation.
- Implemented in **TensorFlow Keras**.
- Final layer uses `sigmoid` activation for **binary segmentation**.

### Loss Function
- **Weighted Focal Loss**  
  Used to address the class imbalance between foreground (LA) and background.
  ```python
  loss = alpha * (1 - pt)^γ * BCE
  ```

### Metrics
- **Dice Score**
- **95th percentile Hausdorff Distance (HD95)**

### Class Weights
- Computed based on pixel counts in the training masks.
- Helps balance the loss function for underrepresented structures like the LA.
